import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import os
from tqdm import trange
import shutil
from sklearn.linear_model import LinearRegression
import random

def phase_shiffting(fid_data):
    phase = np.ones(fid_data.shape)
    phase[1:fid_data.shape[0]:2, :] = np.exp(1j * np.pi)
    fid_data_phase = fid_data * phase
    return fid_data_phase



def fid_to_img(fid_data, mask, method='T1_max'):



    if method == 't2':
        cs_fid = np.multiply(mask, fid_data)
        cs_img = abs(np.fft.fftshift(np.fft.ifftn(np.fft.ifftshift(cs_fid))))
        cs_img_max_value = np.max(cs_img)
        cs_img = cs_img / cs_img_max_value

        img = abs(np.fft.fftshift(np.fft.ifftn(np.fft.ifftshift(fid_data))))
        img_max_value = np.max(img)
        img = img / img_max_value


    elif method == 't1':
        # fid_data = phase_shiffting(fid_data)

        cs_fid = np.multiply(mask, fid_data)

        cs_img = abs(np.fft.fftshift(np.fft.ifftn(np.fft.ifftshift(cs_fid))))
        cs_img_max_value = np.max(cs_img)
        cs_img = cs_img / cs_img_max_value

        img = abs(np.fft.fftshift(np.fft.ifftn(np.fft.ifftshift(fid_data))))
        img_max_value = np.max(img)
        img = img / img_max_value



    return img, cs_img, fid_data, cs_fid, img_max_value, cs_img_max_value

def img_cutting(img, shape):
    img = img.squeeze()
    row = img.shape[0]
    col = img.shape[1]
    new_img = np.zeros([shape, shape], dtype=np.complex128)
    if row == shape:
        diff_col = int((col - shape) / 2)
        new_img[:, :] = img[:, diff_col : diff_col+shape]
    elif col == shape:
        diff_row = int((row - shape) / 2)
        new_img[:, :] = img[diff_row : diff_row+shape, :]
    else:
        diff_row = int((row - shape) / 2)
        diff_col = int((col - shape) / 2)
        new_img[:, :] = img[diff_row: diff_row + shape, diff_col : diff_col+shape]

    return new_img

def gen_mask_gaussian2(kspace_shape, dividend=32, accel_factor=2, seed=None):
    # inspired by https://github.com/facebookresearch/fastMRI/blob/master/common/subsample.py
    shape = kspace_shape
    num_cols = shape[-2]

    center_fraction = (dividend // accel_factor) / 100      #af=2,32 af=3,48
    acceleration = accel_factor

    # Create the mask

    num_low_freqs = int(round(num_cols * center_fraction))
    _, remainder = divmod(num_low_freqs, 2)
    num_low_freqs = num_low_freqs if (remainder == 0) else (num_low_freqs + 1)
    num_high_freqs = int(round(num_cols / acceleration - num_low_freqs))

    high_idx_over = np.round(
        np.random.default_rng(seed).normal(
            num_cols // 2,
            num_cols // 4,
            int(num_high_freqs * 2.5)
        )
    ).astype(int)

    high_idx_over = np.array(list((set(high_idx_over.squeeze().tolist()))))   #set 删除重复数据

    high_idx_over = np.delete(
        high_idx_over, np.where(high_idx_over < 0)
    )

    high_idx_over = np.delete(
        high_idx_over, np.where(high_idx_over >= num_cols)
    )

    low = high_idx_over >= (num_cols // 2 - num_low_freqs // 2)
    high = high_idx_over < (num_cols // 2 + num_low_freqs // 2)
    high_idx_over = np.delete(
        high_idx_over, np.where(low & high)
    )

    if len(high_idx_over) <= num_high_freqs:
        high_idx_tmp = high_idx_over
    else:
        len_over = len(high_idx_over) - num_high_freqs
        random.seed(seed)
        idx_over = random.sample(range(1, len(high_idx_over)), len_over)
        high_idx_tmp = np.delete(high_idx_over, idx_over)

    high_idx = high_idx_tmp
    mask = np.zeros(num_cols)
    pad = (num_cols - num_low_freqs + 1) // 2
    mask[pad: pad + num_low_freqs] = True
    mask[high_idx] = True

    # Reshape the mask
    mask_shape = [1 for _ in shape]
    mask_shape[-1] = num_cols
    mask = mask.reshape(*mask_shape)
    mask = np.reshape(mask, [mask.shape[1], mask.shape[0]])
    fourier_mask = np.repeat(mask.astype(np.float), shape[1], axis=1)

    mask_low = np.zeros(num_cols)
    mask_low[pad: pad + num_low_freqs] = True
    mask_low = mask_low.reshape(*mask_shape)
    mask_low = np.reshape(mask_low, [mask_low.shape[1], mask_low.shape[0]])
    fourier_mask_low = np.repeat(mask_low.astype(np.float), shape[1], axis=1)
    return fourier_mask, fourier_mask_low



accelerations = [4,8,16]

for acceleration in accelerations:
    # kspace_shape = [256,256]
    # mask = np.zeros(kspace_shape)
    # print (mask.shape,kspace_shape)
    # start = mask.shape[0]//2 - mask.shape[0]//(2* acceleration )
    # end = mask.shape[0]//2 + mask.shape[0]//(2* acceleration )
    # mask[start:end, :] = 1
    # print (start,end)
    mask, mask_low = gen_mask_gaussian2([256, 256], dividend=32, accel_factor=acceleration, seed=1108)
    for fid_filename in os.listdir('dataset/fid'):
        fid_path = 'dataset/fid/'+fid_filename

        if not os.path.exists (f'data/t1_fid/'):
            os.mkdir (f'data/t1_fid/')
        if not os.path.exists (f'data/t2_fid/'):
            os.mkdir (f'data/t2_fid/')
        if not os.path.exists (f'data_{acceleration }/'):
            os.mkdir(f'data_{acceleration }/')
        if not os.path.exists (f'data_{acceleration }/'+ fid_filename.split('_')[0]):
            os.mkdir(f'data_{acceleration }/'+ fid_filename.split('_')[0])
        if not os.path.exists (f'data_{acceleration }/'+ fid_filename.split('_')[0]+'_mask'):
            os.mkdir(f'data_{acceleration }/'+ fid_filename.split('_')[0]+'_mask')
        # clean and rebuild dirs
        fid_save_path = f'data/'+fid_filename.split('_')[0]+'_fid' 
        image_save_path = f'data_{acceleration }/'+ fid_filename.split('_')[0]
        masked_image_save_path = f'data_{acceleration }/'+ fid_filename.split('_')[0]+'_mask'
        shutil.rmtree(image_save_path)
        shutil.rmtree(masked_image_save_path)
        os.mkdir(image_save_path)
        os.mkdir(masked_image_save_path)

        fid_data = np.load(fid_path)

        # print (fid_data.shape)
        for index in trange (fid_data.shape[0]) :
            raw_data = fid_data[index]
            for slice in range(raw_data.shape[0]):
                
                
                raw_data_slice = raw_data[slice,: ,:, :]
                # 270,272只留中间
                

                for channel in range(raw_data_slice.shape[0]):
                    cropped_fid = img_cutting(raw_data_slice[channel],256)
                    img, cs_img, _, _, _, _ = fid_to_img(cropped_fid, mask,fid_filename.split('_')[0])
                    fid_save = cropped_fid
                    im_mask  = cs_img[:,:]
                    im_full  = img[:,:]
                    fid_name = fid_save_path + '/%02d_%02d_%02d'%(index,slice,channel)
                    image_name = image_save_path + '/%02d_%02d_%02d'%(index,slice,channel)
                    im_mask_name = masked_image_save_path + '/%02d_%02d_%02d_mask'%(index,slice,channel)
                    np.save(image_name,im_full)
                    np.save(im_mask_name,im_mask)
                    np.save(fid_name,fid_save)
        image = Image.fromarray((img*255).astype('uint8'))
        image.save('t1.jpg')
        image = Image.fromarray((img*255).astype('uint8'))
        image.save('t1_mask.jpg')

                    
        
